Add `dims` support to `sort`, `sort!`, `sortperm` and `sortperm!` by shreyas-omkar · Pull Request #117 · JuliaGPU/AcceleratedKernels.jl

shreyas-omkar · 2026-08-24T09:38:24Z

Adds a dims keyword to the four sorting entry points, so each 1D slice of an array along one dimension is sorted (or permuted) independently, like Base.sort!(A; dims):

import AcceleratedKernels as AK
using CUDA

A = CuArray(rand(Float32, 1000, 1000))
AK.sort!(A; dims=1)                     # every column sorted
B = AK.sort(A; dims=2, rev=true)        # every row sorted, descending
ix = AK.sortperm(A; dims=2)             # linear indices: A[ix] has every row sorted
AK.sortperm!(ix, A; dims=1, by=abs)     # fills ix in place; A is unchanged

The default stays dims=:, which sorts the whole array as one vector while preserving its shape. Integer dims must be in 1:ndims(A); otherwise an ArgumentError is thrown. The lt, by, rev and order keywords apply within each slice. sortperm returns stable linear indices into A; with integer dims, sortperm! requires ix to have the same axes as A, or throws an ArgumentError.

On GPU, block_size controls the tile size and temp supplies an array-sized merge buffer (values for sort!, indices for sortperm!); key buffers may still be allocated. The CPU slice path ignores both keywords. alg=AK.RadixSort() with integer dims throws an ArgumentError. This closes #59 and provides the sorting support needed by JuliaGPU/AMDGPU.jl#1033 and JuliaGPU/GPUArrays.jl#608.

How it works

The GPU merge sort kernels become slice-aware instead of gaining a separate batched sort. A small SliceLayout (src/sort/slices.jl) describes the slices of an array along a dimension as len elements stride apart; every block derives its slice from its linear block index and reads and writes it through a SliceView using a linear offset and stride. The merge logic is unchanged:

Block phase: each block sorts a tile of one slice in local memory, so slices of up to 2 * block_size elements (512 by default) need no global merges. For sort! with by=identity, nontrivial slices take a single launch with no temporary array.
Global phase: the merge passes stay within a slice, so a slice of length n needs ceil(log2(n / tile)) global passes when n > tile, instead of using the total array length N to determine the pass count.

The whole-array sort is a FlatLayout with one slice, for which the slice accessor returns the array itself, retaining direct array indexing. sortperm reuses the key/value merge sort with the linear indices as values, exactly as the flat case does, and the alg=AK.MergeSort(lowmem=true) variant sorts indices while comparing values in the original array. sort! still hoists nonidentity by transforms into keys = by.(v); sortperm evaluates by during comparisons.

On CPU backends each slice is sorted with Base.sort! on a view, with slices spread over tasks; sortperm follows Base and sorts the linear indices of each slice by the values they point to.

Performance

Float32, milliseconds, whole call including allocation, against the backend package's own sort!(A; dims) / sortperm(A; dims) where it has one.

NVIDIA RTX 5080 (compared with CUDA.jl):

size	dims	AK sort!	CUDA.jl sort!	AK sortperm	CUDA.jl sortperm
1024 x 1024	1	0.16	0.31	0.19	0.84
1024 x 1024	2	0.18	0.42	0.25	1.02
256 x 4096	1	0.17	0.38	0.21	0.91
256 x 4096	2	0.26	0.64	0.38	1.34
4096 x 256	1	0.21	0.44	0.24	0.95
4096 x 256	2	0.17	0.49	0.22	1.10
64 x 16384	1	0.38	0.89	0.48	1.52
16384 x 64	1	0.27	0.63	0.30	1.20
128 x 128 x 64	3	0.39	1.30	0.49	2.14

For reference, CUB's DeviceSegmentedSort (what torch.sort uses) takes 0.14 ms for SortKeys and 0.13 ms for SortPairs on the contiguous 1024 x 1024 case, and 0.04 / 0.14 ms for 256 slices of 4096 (kernel time only, excluding allocation). It pairs a radix sort with size-specialised kernels for small segments, is not stable, and only handles contiguous slices.

Apple M1 (compared with Metal.jl):

size	dims	AK sort!	Metal.jl sort!	AK sortperm	Metal.jl sortperm
1024 x 1024	1	4.1	8.7	7.5	9.0
1024 x 1024	2	4.1	8.9	7.3	9.5
256 x 4096	1	4.1	30.8	6.2	31.4
4096 x 256	1	6.1	3.3	8.3	3.5
64 x 16384	1	11.8	119.0	16.6	124.1
16384 x 64	1	8.3	2.6	10.6	3.4
128 x 128 x 64	3	11.9	121.3	16.7	125.2

Intel Xe iGPU (oneAPI.jl forwards sort! to AK and its sortperm(A; dims) errors, so no vendor column):

size	dims	AK sort!	AK sortperm
1024 x 1024	1	6.3	10.9
1024 x 1024	2	6.6	11.1
256 x 4096	1	7.9	10.1
4096 x 256	1	8.8	14.0
64 x 16384	1	23.1	24.3
16384 x 64	1	11.7	17.9

Not done here

Slices much smaller than a block tile still get a whole block each, so very short slices (tens of elements) leave most threads idle.
Slices with stride greater than one are accessed directly; there is no transpose-based path to improve coalescing.
RadixSort along dims would need a segmented radix sort.

shreyas-omkar · 2026-08-25T11:09:56Z

@christiangnrd Please take a look. :)

shreyas-omkar · 2026-08-26T07:04:29Z

NVIDIA RTX 5080 (CUDA 13.3.0)

N-D sort(A; dims) Float32, 1,048,576 elements

size	dims	AK ms	CUDA ms	AK/CUDA
1024 x 1024	1	0.649	0.349	1.86x
1024 x 1024	2	0.671	0.469	1.43x
256 x 4096	1	0.641	0.435	1.47x
256 x 4096	2	0.682	0.706	0.97x
4096 x 256	1	0.659	0.494	1.33x
4096 x 256	2	0.663	0.543	1.22x

christiangnrd · 2026-09-03T15:43:32Z

Ignoring the benchmarks that this PR doesn't touch, these numbers don't look great

Sort or permute each 1D slice along an integer `dims` independently, matching Base. The default `dims=:` keeps the existing flat behaviour. There is no batched sort kernel, so each element is tagged with its slice and the whole array is sorted once by (slice, value), then scattered back into place. sortperm carries the original index as the payload and uses it to break ties, keeping the permutation stable. This reuses the backend's tuned sort and runs unchanged on CPU and every GPU backend.

Replace the tag-and-sort implementation of `dims` (tuple keys, a global sort of the whole array and a scatter back) with merge sort kernels that know about slices. A `SliceLayout` describes the 1D slices of an array along a dimension as `len` elements `stride` apart; each block maps its linear index to a slice and a tile within it, and accesses the slice through a `SliceView`. Slices up to twice the block size are sorted in local memory in one launch, larger ones get per-slice global merge passes. The whole-array sort is a one-slice `FlatLayout` whose accessor returns the array itself, so that path is unchanged. `merge_sort!`, `merge_sort_by_key!`, `merge_sortperm!` and `merge_sortperm_lowmem!` take `dims` directly; `sortperm` keeps linear indices as the values of a key/value sort, which is stable. On CPU each slice is sorted with `Base.sort!` on a view, slices spread over tasks, and `sortperm` sorts linear indices with a `Perm` ordering like Base. `RadixSort` with `dims` throws instead of silently switching algorithms, and the low-level sortperm entry points now default to `lt=isless`. Sorting 1024x1024 Float32 along a dimension on an RTX 5080 drops from 0.51 to 0.16 ms (CUDA.jl's bitonic sort: 0.31 ms), sortperm from 0.66 to 0.19 ms; the M1 and an Intel iGPU see 3-7x. The dims tests now also run on the CPU path and cover tile boundaries, global merge passes, ties, NaNs, strided views and empty slices.

maleadt · 2026-09-16T10:47:50Z

I reworked the implementation on top of your commit, keeping the dims API and most of the tests.

Tag-and-sort (the same scheme JuliaGPU/AMDGPU.jl#1033 uses as a stopgap) is what made the numbers disappointing: it allocates an array of (slice, value) tuples ((slice, value, index) for sortperm) plus a merge buffer of that size, adds tag and scatter passes, compares tuples in the binary-search hot loop, and runs log2(N) merge passes over the whole array where each slice only needs log2(n).

The existing merge kernels now use a SliceLayout and SliceView to sort each slice independently. Slices up to 2 * block_size need only the local-memory sort phase; larger slices get global merges whose pass count depends on slice length, not total array length. sortperm uses the stable key/value merge sort, while alg=AK.MergeSort(lowmem=true) sorts indices by looking up the original values. sort! retains its by hoisting.

The flat layout still indexes the array directly. Whole-array timings remain close: 16M Float32 on the RTX 5080 takes 9.0 vs 9.1 ms; M1 and Intel iGPU timings are in the PR description.

Other changes:

RadixSort with integer dims now throws ArgumentError instead of silently falling back, and GPU slice sorts honour temp. The CPU slice path ignores temp.
The CPU path uses Base.sort! on views, distributing slices over tasks. sortperm sorts linear indices with Base's Perm ordering.
The dims testsets now also run on the threaded CPU path. Added coverage includes tile boundaries and global merges with duplicates (stability), temp/block_size/lowmem, NaNs and signed zeros, 4D value sorts, and empty slices. - Docs example in api/sort.md.

Performance: Float32, ms, whole call including allocation. "Before" is this PR as submitted, "after" this rework, "vendor" the backend package's own sort!(A; dims) / sortperm(A; dims).

Device	size	dims	sort! before	sort! after	vendor sort!	sortperm before	sortperm after	vendor sortperm
RTX 5080	1024 x 1024	1	0.51	0.16	0.31	0.66	0.19	0.84
RTX 5080	1024 x 1024	2	0.53	0.18	0.42	0.68	0.25	1.02
RTX 5080	256 x 4096	2	0.54	0.26	0.64	0.70	0.38	1.34
RTX 5080	16384 x 64	1	0.53	0.27	0.63	0.69	0.30	1.20
RTX 5080	128 x 128 x 64	3	0.52	0.39	1.30	0.66	0.49	2.14
Apple M1	1024 x 1024	1	18.5	4.1	8.7	23.6	7.5	9.0
Apple M1	256 x 4096	1	18.6	4.1	30.8	23.6	6.2	31.4
Apple M1	16384 x 64	1	18.3	8.3	2.6	23.6	10.6	3.4
Intel Xe iGPU	1024 x 1024	1	43.7	6.3	-	64.3	10.9	-
Intel Xe iGPU	16384 x 64	1	44.0	11.7	-	65.6	17.9	-

oneAPI.jl forwards sort! to AK, so it has no independent number, and its sortperm(A; dims) errors. CUB's DeviceSegmentedSort::SortKeys does the contiguous 1024 x 1024 case in 0.14 ms on the RTX 5080, so there is still headroom for small slices, which CUB handles with warp-level sorts; that seemed out of scope here.

Sort the network ascending in `ord` with the reflected first step of every merge level, and skip comparators whose partner lies past the end: any length sorts without padding, so the sentinel buffer, the eltype whitelist and the Forward/Reverse restriction go away and NaNs order like Base. Every kernel takes the slice layout from JuliaGPU#117, so `dims` needs no separate path, slices of any length work, and slices shorter than a workgroup share a tile. Tunables live on `BitonicSort(; block_size, items_per_thread)` with backend defaults from `bitonic_defaults`.

luraess · 2026-09-16T13:37:05Z

I ran this on AMD, since it was not in the validation set above. RX 7900 XTX, ROCm 6.4.4, Julia 1.12.

Your sort_dims and sortperm_dims testsets pass on ROCBackend(): 1584 and 3123 tests, no failures.

Float32, milliseconds, minimum of 20 with a warm-up, timing the call itself with a fresh input each repetition. The comparison column is the tag-and-sort workaround this replaces in JuliaGPU/AMDGPU.jl#1033, where every element was tagged with its slice index and the array sorted once.

size	dims	AK sort!	workaround	AK sortperm	workaround
1024 x 1024	1	0.49	1.03	0.52	2.21
1024 x 1024	2	0.51	0.86	0.48	2.43
256 x 4096	1	0.54	0.86	0.32	1.78
256 x 4096	2	0.71	0.81	1.83	1.85
4096 x 256	1	0.57	0.86	0.40	1.78
4096 x 256	2	0.33	0.84	0.31	1.82
64 x 16384	1	0.91	0.77	0.73	2.00
16384 x 64	1	0.71	0.89	0.52	1.72
128 x 128 x 64	3	0.87	0.96	0.96	2.12

8 of 9 cases faster for both, up to 2.5x for sort! and 5.9x for sortperm. The two exceptions are exactly your "Not done here" items: 64 x 16384, dims=1 is 16384 slices of 64 elements each taking a whole block tile, and 256 x 4096, dims=2 is the strided case.

No regression on the flat path: 1M Float32 sort! is 0.75 ms on both 0.4.3 and main.

One observation: these land around 3x the RTX 5080 column, on a GPU of comparable memory bandwidth, so there may be ROCm-side headroom. Measurement details differ, so take it as a hint rather than a number.

A tagged release would be much appreciated. AMDGPU.jl needs two lines on top of this: JuliaGPU/AMDGPU.jl#1076

maleadt · 2026-09-16T13:56:50Z

Nice, thanks! I do want to first do a full pass over the new abstractions before releasing though, so that will still take a bit. Also, the plan is to relatively quickly wire this up to GPUArrays.jl, so maybe it's unnecessary to put this in AMDGPU.jl at all?

luraess · 2026-09-16T14:24:10Z

Also, the plan is to relatively quickly wire this up to GPUArrays.jl, so maybe it's unnecessary to put this in AMDGPU.jl at all?

I have this ready now JuliaGPU/AMDGPU.jl#1076 but yeah - if the GPUArray path would be better and land relatively soon then we could just wait on that. Else tagging a minor of AK would be helpful.

maleadt · 2026-09-16T16:13:12Z

Why not land the slow JuliaGPU/AMDGPU.jl#1033 as a stopgap and switch over to the fast implementation in a couple of weeks?

shreyas-omkar force-pushed the sh/sort-dims branch from 8410490 to b7b31ad Compare August 24, 2026 09:47

shreyas-omkar changed the title ~~Add dims support to sort and sort!~~ Add dims support to sort, sort!, sortperm and sortperm! Aug 24, 2026

shreyas-omkar marked this pull request as draft August 24, 2026 09:49

shreyas-omkar marked this pull request as ready for review August 26, 2026 06:51

This was referenced Sep 8, 2026

Segmented radix for sort(A; dims) via onesweep #120

Closed

perf(sort): packed-key fast path for sort(A; dims=1) #121

Closed

shreyas-omkar force-pushed the sh/sort-dims branch from b7b31ad to aaed884 Compare September 16, 2026 09:19

maleadt self-assigned this Sep 16, 2026

maleadt changed the title ~~Add dims support to sort, sort!, sortperm and sortperm!~~ Add dims support to sort, sort!, sortperm and sortperm! Sep 16, 2026

maleadt merged commit 017520f into JuliaGPU:main Sep 16, 2026
53 of 54 checks passed

maleadt mentioned this pull request Sep 16, 2026

Add BitonicSort, a GPU sorting network for small arrays and short slices #126

Merged

luraess mentioned this pull request Sep 16, 2026

Support dims in sort and sortperm, on top of AK JuliaGPU/AMDGPU.jl#1076

Draft

luraess mentioned this pull request Sep 16, 2026

Support dims in sort, sort!, sortperm and sortperm! JuliaGPU/AMDGPU.jl#1033

Open

Navigation Menu

Add `dims` support to `sort`, `sort!`, `sortperm` and `sortperm!` - #117

Add `dims` support to `sort`, `sort!`, `sortperm` and `sortperm!`#117
maleadt merged 2 commits into
JuliaGPU:mainfrom
shreyas-omkar:sh/sort-dims

shreyas-omkar commented Aug 24, 2026 •

edited by maleadt

Loading

Uh oh!

shreyas-omkar commented Aug 25, 2026

Uh oh!

shreyas-omkar commented Aug 26, 2026 •

edited by christiangnrd

Loading

Uh oh!

christiangnrd commented Sep 3, 2026

Uh oh!

maleadt commented Sep 16, 2026 •

edited

Loading

Uh oh!

Uh oh!

luraess commented Sep 16, 2026

Uh oh!

maleadt commented Sep 16, 2026

Uh oh!

luraess commented Sep 16, 2026

Uh oh!

maleadt commented Sep 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

shreyas-omkar commented Aug 24, 2026 • edited by maleadt Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

How it works

Performance

Not done here

Uh oh!

shreyas-omkar commented Aug 25, 2026

Uh oh!

shreyas-omkar commented Aug 26, 2026 • edited by christiangnrd Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

NVIDIA RTX 5080 (CUDA 13.3.0)

Uh oh!

christiangnrd commented Sep 3, 2026

Uh oh!

maleadt commented Sep 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

luraess commented Sep 16, 2026

Uh oh!

maleadt commented Sep 16, 2026

Uh oh!

luraess commented Sep 16, 2026

Uh oh!

maleadt commented Sep 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

shreyas-omkar commented Aug 24, 2026 •

edited by maleadt

Loading

shreyas-omkar commented Aug 26, 2026 •

edited by christiangnrd

Loading

maleadt commented Sep 16, 2026 •

edited

Loading